/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.
 *
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 */
#ifndef ATVC_REDUCE_TILING_H
#define ATVC_REDUCE_TILING_H
#include "common/atvc_opdef.h"
#include "reduce/tiling/reduce_op_tiling.h"
#include "reduce/tiling/tiling_def.h"

namespace ATVC {
namespace Host {

/**
 * @brief 计算Reduce的TilingData和策略参数
 * @param inputShape 输入张量的形状。
 * @param reduceDim 需要进行Reduce操作的具体维度。
 * @param policy 输出参数。
 * @param param 输出参数。
 * @return bool 返回true表示计算成功，false表示失败。
 */
template<class OpTraits>
bool CalcReduceTiling(std::vector<int64_t> inputShape, std::vector<int64_t> reduceDim, ReducePolicy* policy, ReduceParam* param)
{
    struct ReduceTilingHyperParam {
        int32_t basicBlock = 16 * 1024;  // 最大为UB的总大小的1/3
        int32_t nBufferNum = 2;          // 每个Queue中的Tensor数量
    };
    using inputDTypeList = typename OpTraits::In::types;
    static constexpr size_t inTensorSumBytes = ATVC::TypeListReduce<inputDTypeList, SizeValue<0>, SumSizes>::Type::value;
    if (inTensorSumBytes == 0) {
        printf("[ERROR] Tiling Error: OpTraits Input cannot be null!\n");
        return false;
    }
    using DataType = typename ATVC::TypeListGet<inputDTypeList, 0>::Type;
    auto inputDtype = GetOriInputType<DataType>();
    // check output
    using outputDTypeList = typename OpTraits::Out::types;
    using DataTypeOut = typename ATVC::TypeListGet<outputDTypeList, 0>::Type;
    if (GetOriInputType<DataTypeOut>() != inputDtype) {
        printf("[ERROR] Output Dtype is different with input's!\n");
        return false;
    }
    ReduceTilingInputParam opInput = {reduceDim, inputShape, inputDtype, GetPromoteDataType(inputDtype)};
    OpTiling::ReduceOpTiling tiling(opInput, policy, param);
    if(tiling.Run() != 0) {
        printf("[ERROR] Tiling Error\n");
        return false;
    }
    printf("workspaceSize %d\n", param->workspaceSize);
    printf("policy  (%d, %d, %d)\n", policy->patternID, policy->loopARCount, policy->loopInnerARCount);
    printf("tilingData.factorACntPerCore %zu\n", param->tilingData.factorACntPerCore);
    printf("tilingData.factorATotalCnt %zu\n", param->tilingData.factorATotalCnt);
    printf("tilingData.ubFactorA %zu\n", param->tilingData.ubFactorA);
    printf("tilingData.factorRCntPerCore %zu\n", param->tilingData.factorRCntPerCore);
    printf("tilingData.factorRTotalCnt %zu\n", param->tilingData.factorRTotalCnt);
    printf("tilingData.ubFactorR %zu\n", param->tilingData.ubFactorR);
    printf("tilingData.groupR %zu\n", param->tilingData.groupR);
    printf("tilingData.outSize %zu\n", param->tilingData.outSize);
    printf("tilingData.basicBlock %zu\n", param->tilingData.basicBlock);
    printf("tilingData.coreNum %d\n", param->tilingData.coreNum);
    printf("nBufferNum %d\n", param->nBufferNum);
    return true;
};
} // Host
} // ATVC
#endif